#include "kiva/node.h"

#include <glog/logging.h>

namespace py = pybind11;
namespace kiva {

namespace {
template <NodeType first, NodeType last>
constexpr bool NodeTypeCheck(NodeType val) {
  static_assert(int(first) <= int(last));
  return int(first) <= int(val) && int(val) <= int(last);
}
}  // namespace

bool Node::IsConstant(NodeType type) {
  return type == NodeType::kConstant ||
         NodeTypeCheck<NodeType::kConstantFirst, NodeType::kConstantLast>(type);
}

bool Node::IsBinOp(NodeType type) {
  return NodeTypeCheck<NodeType::kBinaryExprFirst, NodeType::kBinaryExprLast>(
      type);
}

bool Node::IsUnaryOp(NodeType type) {
  return NodeTypeCheck<NodeType::kUnaryExprFirst, NodeType::kUnaryExprLast>(
      type);
}

bool Node::IsBoolOp(NodeType type) {
  return NodeTypeCheck<NodeType::kBoolOpFirst, NodeType::kBoolOpLast>(type);
}

bool Node::IsCompare(NodeType type) {
  return NodeTypeCheck<NodeType::kCompareFirst, NodeType::kCompareLast>(type);
}

bool Node::IsExpression(NodeType type) {
  return NodeTypeCheck<NodeType::kExpressionFirst, NodeType::kExpressionLast>(
      type);
}

bool Node::IsStatement(NodeType type) {
  return NodeTypeCheck<NodeType::kStatementFirst, NodeType::kStatementLast>(
      type);
}

bool Node::IsStatementNoChild(NodeType type) {
  return NodeTypeCheck<NodeType::kStatementNoChildFirst,
                       NodeType::kStatementNoChildLast>(type);
}

bool Node::HasNoChild(NodeType type) {
  return IsStatementNoChild(type) || IsConstant(type);
}

Node::Node(NodeType type, std::initializer_list<NodeRef> children)
    : type_(type), value_(ListType()) {
  auto& list = as_list();
  list.reserve(children.size());
  for (NodeRef child : children)
    list.push_back(std::unique_ptr<Node>(child.get()));
}

Node::Node(ObjectType val)
    : type_(NodeType::kConstant), value_(std::move(val)) {}

Node::Node(std::string_view val)
    : type_(NodeType::kName), value_(StringType(val)) {}

Node::Node(LongType val) : type_(NodeType::kLong), value_(val) {}

void Node::set_type(NodeType type) {
  // check that the new type is compatible with the old type.
  DCHECK(HasNoChild(type_) == HasNoChild(type));
  type_ = type;
}

py::object Node::object_value() const {
  DCHECK(type() == NodeType::kConstant);
  return std::get<ObjectType>(value_);
}

long Node::long_value() const {
  DCHECK(type() == NodeType::kLong);
  return std::get<LongType>(value_);
}

const Node::StringType& Node::string_value() const {
  DCHECK(type() == NodeType::kName);
  return std::get<StringType>(value_);
}

Node::ListType& Node::as_list() {
  return std::get<ListType>(value_);
}

const Node::ListType& Node::as_list() const {
  return std::get<ListType>(value_);
}

// Take ownership of this node.
void Node::AddChild(NodeRef node) {
  AddChild(std::unique_ptr<Node>(node.get()));
}

void Node::AddChild(std::unique_ptr<Node> node) {
  as_list().push_back(std::move(node));
}

NodeRef Node::NewWithChildren(NodeType type,
                              std::initializer_list<NodeRef> children) {
  return new Node(type, children);
}

NodeRef Node::NewWithChildren(NodeType type, NodeRef child) {
  return new Node(type, {child});
}

NodeRef Node::NewName(std::string_view val) {
  return new Node(val);
}

NodeRef Node::NewNodeType(NodeType val) {
  return NewLong((long)val);
}

NodeRef Node::NewLong(LongType val) {
  return new Node(val);
}

NodeRef Node::NewConstant(long val) {
  return new Node(py::int_(val));
}

NodeRef Node::NewConstant(unsigned long val) {
  return new Node(py::int_(val));
}

NodeRef Node::NewConstant(double val) {
  return new Node(py::float_(val));
}

NodeRef Node::NewConstant(bool val) {
  return new Node(py::bool_(val));
}

NodeRef Node::NewConstant(std::string_view val) {
  return new Node(py::str(val.data(), val.size()));
}

NodeRef Node::NewBinOp(NodeRef lhs, NodeType type, NodeRef rhs) {
  DCHECK(Node::IsBinOp(type));
  return NewWithChildren(type, {lhs, rhs});
}

NodeRef Node::NewAssign(NodeRef lhs, NodeRef rhs) {
  return NewWithChildren(NodeType::kAssign, {lhs, rhs});
}

NodeRef Node::NewAugAssign(NodeRef lhs, NodeType type, NodeRef rhs) {
  DCHECK(Node::IsBinOp(type));
  return NewWithChildren(NodeType::kAssign, {lhs, NewNodeType(type), rhs});
}

NodeRef Node::NewUnaryOp(NodeType type, NodeRef operand) {
  DCHECK(Node::IsUnaryOp(type));
  return NewWithChildren(type, {operand});
}

NodeRef Node::NewBoolOp(NodeRef lhs, NodeType type, NodeRef rhs) {
  DCHECK(Node::IsBoolOp(type));
  return NewWithChildren(type, {lhs, rhs});
}

NodeRef Node::NewCompare(NodeRef lhs, NodeType type, NodeRef rhs) {
  DCHECK(Node::IsCompare(type));
  // If lhs is a compare, add type and rhs to its children.
  // If lhs is not a compare, create a compare and add lhs, type and rhs to
  // its children.
  NodeRef cmp = lhs;
  if (cmp->type() != NodeType::kCompare) {
    cmp = NewWithChildren(NodeType::kCompare, {lhs});
  }
  cmp->AddChild(NewNodeType(type));
  cmp->AddChild(rhs.get());
  return cmp;
}

NodeRef Node::NewAttribute(NodeRef lhs, std::string_view name) {
  // todo: name may be dotted, so split them.
  return NewWithChildren(NodeType::kAttribute, {lhs, NewName(name)});
}

NodeRef Node::NewCall(NodeRef func, NodeRef args) {
  return NewWithChildren(NodeType::kCall, {func, args});
}

NodeRef Node::NewSubscript(NodeRef node, NodeRef index) {
  return NewWithChildren(NodeType::kSubscript, {node, index});
}

NodeRef Node::NewAlias(NodeRef name, std::string_view alias) {
  return NewWithChildren(NodeType::kAlias, {name, NewName(alias)});
}

NodeBuilder::NodeBuilder(NodeType type)
    : node_(Node::NewWithChildren(type).get()) {}

NodeBuilder::NodeBuilder(NodeRef node) : node_(node.get()) {}

NodeRef NodeBuilder::Build() {
  return NodeRef(node_.release());
}

NodeBuilder& NodeBuilder::AddChildren(std::initializer_list<NodeRef> children) {
  DCHECK(node_);
  for (auto child : children)
    node_->AddChild(child.get());
  return *this;
}

NodeBuilder& NodeBuilder::AddChild(NodeRef child) {
  DCHECK_NOTNULL(node_)->AddChild(child.get());
  return *this;
}
}  // namespace kiva
